import torch
import torch.nn as nn
import torch.nn.functional as F

from ecgcmr.imaging.img_models.ResNetModel import generate_model


class ResNet3D_Encoder(nn.Module):
    def __init__(self, model_depth: int = 50, zero_init_residual: bool = True) -> None:
        super().__init__()
        
        self.encoder = generate_model(model_depth=model_depth, zero_init_residual=zero_init_residual)
        self.encoded_dim = self.encoder.encoded_dim
        self.avgpool = nn.AvgPool3d(kernel_size=(1, 3, 3), stride=1, padding=0)

        self.T_contrastive = None

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        encoded_images = self.encoder.forward_features(images) # B, C, T', H', W' (T'= T)
        embeddings = self.avgpool(encoded_images).flatten(start_dim=2).permute(0, 2, 1) # B, T', C

        return embeddings

    def _fill_avg_pool_flatten_interpolate(self, 
                                           encoded_images: torch.Tensor,
                                           delta_steps: list) -> torch.Tensor:
        B = encoded_images.shape[0]
        final_batch = []

        for i in range(B):
            last_frame = encoded_images[i, :, -1, :, :].unsqueeze(1)
            repeat_count = delta_steps[i]

            if repeat_count > 0:    
                repeated_frames = last_frame.repeat(1, repeat_count, 1, 1)
                filled_encoded_image = torch.cat((encoded_images[i], repeated_frames), dim=1)

            else:
                filled_encoded_image = encoded_images[i]

            avg_filled_encoded_image = self.avgpool(filled_encoded_image) # C, T', 1, 1
            flattened_avg_filled_encoded_image = avg_filled_encoded_image.flatten(start_dim=1) # C, T

            interpolated_flattened_avg_filled_encoded_image = F.interpolate(flattened_avg_filled_encoded_image.unsqueeze(0),
                                                                            size=self.T_contrastive,
                                                                            mode='linear', align_corners=None)
            
            interpolated_flattened_avg_filled_encoded_image = torch.squeeze(interpolated_flattened_avg_filled_encoded_image)
            
            final_batch.append(interpolated_flattened_avg_filled_encoded_image)

        final_batch = torch.stack(final_batch, dim=0)

        return final_batch

    def forward_multimodal(self, images: torch.Tensor, delta_steps: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        TODO: Fix global calculation and make more efficient _fill_avg_pool_flatten_interpolate function
        """
        encoded_images = self.encoder.forward_features(images) # B, C, T', H', W'

        local_embeddings = self._fill_avg_pool_flatten_interpolate(encoded_images=encoded_images,
                                                                   delta_steps=delta_steps) # B, C, T_contr
        
        global_embeddings = local_embeddings.mean(dim=-1) # B, C

        return local_embeddings, global_embeddings
 